Kernel regularization: opt-in Cholesky-solve evidence term + relative jitter - #437

Closes leg 2 of PyAutoMind draft/feature/autoarray/regularization_jax_gradient_gaps.md.

Both changes are opt-in and default-off, so default evidence values are unchanged and archived results stay comparable — the constraint the #391 adversarial probe verdict imposed.

Leg 2a — s^T H s via one Cholesky solve

The kernel schemes define H = coefficient * C^-1 and were forming that inverse explicitly. Forming an inverse is slower and less accurate than solving: the error is amplified by cond(C), which reaches ~1e9 on the clustered traced vertices of the kNN mesh families, putting a noise floor on the likelihood the sampler sees.

regularization_term now computes coefficient * s^T C^-1 s from a single Cholesky solve against s. Shape mirrors the log-det shortcut shipped in #391 exactly: a regularization_term_from hook on AbstractRegularization returning None by default, overridden by the four kernel schemes, consulted behind a new Settings.regularization_term_method ("matmul" default, "cho_solve" opt-in).

Measured on a clustered fixture at cond(C) = 3.2e9, graded against an exactly-known reference (choosing s = Cv makes C^-1 s = v, so the true form is s^T v and needs no inverse anywhere):

Why a separate setting from log_det_method

Deliberate. The two evidence terms can be moved onto their exact factorizations independently, which is what makes an evidence shift attributable to one term rather than both. Gating s^T H s on a setting named for log-determinants would also have been a lie in the name.

Scope limit worth stating

This does not remove the explicit inverse from the inversion as a whole. curvature_reg_matrix is a dense F + H feeding the dense solve for the reconstruction, so H is still formed there. Only the evidence terms avoid it. Doing better needs an iterative-solver design — a different piece of work.

Leg 2b — relative jitter

The covariance jitter was a fixed absolute 1e-8 * I, which only makes sense when diag(C) ~ 1. Measured: true for the three unweighted kernels (K(0) == 1), false for MaternAdaptKernel, whose C_ii = w_i^2 spans the adaptive-weight dynamic range.

Distortion of the faintest pixel, 40-pixel fixture:

inner_coefficient/outer_coefficient are free model parameters, so a sampler can reach that bottom row mid-fit — silently, with no exception and no NaN, just wrong smoothing.

jitter_relative=True applies C_ii *= (1 + jitter). Since C = D^(1/2) R D^(1/2), that is exactly D^(1/2) (R + jitter I) D^(1/2) — the jitter lands on the correlation matrix, so every pixel gets the same relative protection whatever its scale.

A rejected alternative, recorded so it isn't retried: jitter = N * eps * max(diag) ("as small as possible above the round-off floor") fixes the distortion but lets cond(C) reach 3.2e15, at the edge of float64 — reintroducing exactly the noise this branch exists to remove. The correlation-relative rule leaves conditioning unchanged (3.16e9 both conventions, measured).

Two traps found while implementing, both pinned by tests

• MaternAdaptKernel passes coefficient=0.0 to MaternKernel.__init__ (its weights live inside C_w). Inheriting the parent's term would have silently zeroed it — wrong but finite, the worst failure mode. It needs its own override.
• GaussianKernel's formed matrix carries a symmetrisation plus trace-scaled jitter that the analytic shortcut excludes — consistent with how its log-det shortcut already behaves, since both exist only to guard a factorization these paths avoid.

Schemes with no factorization (Constant, Adapt, the split families) return None, and one None falls the whole computation back to the formed matrix, so mixed inversions stay correct.

Testing

29 new tests across three files:

• test_kernel_regularization_term.py (9) — shortcut vs formed matrix for all four schemes, the gate in both directions, the fallback, the MaternAdapt zero-coefficient trap, the ill-conditioning gain.
• test_kernel_jitter_relative.py (7) — default byte-identical to previous behaviour, relative jitter as a pure diagonal rescaling, the adaptive-weight bug, conditioning preserved on clustered vertices, flag threaded to all three covariance call sites per scheme.
• test_kernel_jax_gradients.py (13) — the JAX leg of the gate.

JAX certification:

Full suite: 936 passed, 52 skipped. Three test_transformer.py failures are pre-existing on main and unrelated (missing optional pynufft).

Matérn's JAX path is not covered by the new JAX tests — it needs the modified Bessel from tfp-nightly, an optional-of-an-optional. Gaussian and Exponential exercise the same shared apply_jitter / quadratic_form_via_cholesky code without it. The workspace-level autolens_workspace_test/scripts/imaging/jax_grad/regularization.py still covers the Matérn end-to-end and is unaffected by this branch's defaults.

Generated by Claude Code